#!/usr/bin/env python
import rclpy
from rclpy.node import Node
from rclpy.parameter import Parameter
from sensor_msgs.msg import JointState
from so101_hw_interface.motors.feetech import FeetechMotorsBus, OperatingMode
from so101_hw_interface.motors import Motor, MotorNormMode, MotorCalibration
import time
import math
import json
import pathlib
import sys
import time

CALIB_PATH = pathlib.Path.home() / ".calibrate" / "so101_leader_calibrate.json"

class LeaderArmPublisher(Node):
    def __init__(self):
        super().__init__('so101_leader_publisher')

        self.declare_parameter('port', '/dev/ttyACM0')
        self.port = self.get_parameter('port').get_parameter_value().string_value

        self.declare_parameter('calibrate', False)
        self.is_calibrate_mode = self.get_parameter('calibrate').get_parameter_value().bool_value
        if self.is_calibrate_mode:
            self.get_logger().warn(">>> 节点以 [校准模式] 启动 <<<")
        
        self.get_logger().info(f"Leader Publisher: Using port: {self.port}")

        self.leader_state_pub = self.create_publisher(
            JointState,
            '/so101_leader/joint_states',
            10
        )

        # 舵机定义
        self.joints = {
            "1": {"id": 1, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
            "2": {"id": 2, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
            "3": {"id": 3, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
            "4": {"id": 4, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100},
            "5": {"id": 5, "model": "sts3215", "mode": MotorNormMode.RANGE_M100_100}, # "5" 是全周电机 (Wrist Roll)
            "6": {"id": 6, "model": "sts3215", "mode": MotorNormMode.RANGE_0_100},
        }
        self.joint_names = list(self.joints.keys()) # ["1", "2", "3", "4", "5", "6"]

        self._home_offsets: dict[str, int] | None = None
        self._rad_per_step = (2 * math.pi) / 4096.0 # 转换系数
        # --- 存储 Leader "Home" 姿态的真实弧度偏移量 ---
        self._reset_rad_offsets: dict[str, float] = {}
        self.calibration_data: dict[str, MotorCalibration] | None = None 
        self.bus_ = None

        # 连接总线
        try:
            motors = {
                name: Motor(params["id"], params["model"], params["mode"])
                for name, params in self.joints.items()
            }
            self.bus_ = FeetechMotorsBus(self.port, motors)
            self.bus_.connect()
            self.get_logger().info(f"Connected to motor bus on {self.port}")
        except Exception as e:
            self.get_logger().error(f"Failed to connect to bus on {self.port}: {e}")
            self.bus_ = None 
            rclpy.shutdown() 
            return

        if not self._load_or_run_calibration():
            # 如果校准失败 (无论是自动加载还是手动), 都关闭节点
            self.get_logger().error("校准/加载失败。正在关闭...")
            if self.bus_:
                self.bus_.disconnect()
            rclpy.shutdown()
            return
        
        # 从 MotorCalibration 对象中提取 home offsets
        self._home_offsets = {name: cal.homing_offset for name, cal in self.calibration_data.items()}
        self.get_logger().info(f"Leader 臂已加载 home offsets: {self._home_offsets}")

        # 计算 "Home" 位置的真实弧度偏移量
        self._calculate_reset_offsets()

        # 创建定时器
        timer_period = 0.02  # 50hz
        self.timer = self.create_timer(timer_period, self.timer_callback)
        self.get_logger().info("✅ Leader arm publisher 启动成功, 正在读取舵机...")

    # 计算复位偏移量的方法
    def _calculate_reset_offsets(self):
        """
        使用【硬编码的原始步数】('reset_raw_pose')
        来计算 'Home' 姿态的真实弧度偏移量。
        """
        self.get_logger().info("正在计算 Leader 臂的复位弧度偏移量...")
        if self.calibration_data is None:
            self.get_logger().error("无法计算偏移量: 校准数据未加载。")
            return

        reset_raw_pose = {
            '1': 2060, 
            '2': 884,
            '3': 3095,
            '4': 903,
            '5': 879,
            '6': 1595
        }
        self.get_logger().info(f"使用硬编码的 'Home' 姿态 (原始步数): {reset_raw_pose}")

        for name in self.joint_names:
            if name not in self.calibration_data or name not in reset_raw_pose:
                self.get_logger().warn(f"跳过关节 {name}: 在校准或'reset_raw_pose'中未找到。")
                self._reset_rad_offsets[name] = 0.0 # 默认为 0 偏移
                continue
            
            cal = self.calibration_data[name]
            home = cal.homing_offset
            # 直接使用正确的原始步数
            reset_raw = reset_raw_pose[name]
            
            # (真实弧度) = (原始步数 - 零点步数) * 弧度/步数
            reset_rad = (reset_raw - home) * self._rad_per_step
            self._reset_rad_offsets[name] = reset_rad
        
        self.get_logger().info("成功计算了 Leader 臂的复位弧度偏移量:")
        for name, offset in self._reset_rad_offsets.items():
             self.get_logger().info(f"  - 关节 {name}: {offset:.2f} rad")

    def _save_calibration(self, calibration_data: dict[str, MotorCalibration]):
        """将完整的校准数据保存到 JSON 文件"""
        
        # 必须将 MotorCalibration 对象转换回 dict 以便 JSON 序列化
        data_to_save = {
            name: {
                "id": cal.id,
                "drive_mode": cal.drive_mode,
                "homing_offset": cal.homing_offset,
                "range_min": cal.range_min,
                "range_max": cal.range_max,
            }
            for name, cal in calibration_data.items()
        }
        
        try:
            CALIB_PATH.parent.mkdir(parents=True, exist_ok=True)
            with open(CALIB_PATH, 'w') as f:
                json.dump(data_to_save, f, indent=4) # 保存 dict 版本
            self.get_logger().info(f"校准数据已保存到: {CALIB_PATH}")
        except Exception as e:
            self.get_logger().error(f"无法保存校准文件: {e}")

    def _run_calibration(self) -> dict[str, MotorCalibration] | None:
        """运行交互式校准程序以查找 Homing Offsets 和 Ranges"""
        self.get_logger().info("\n--- 开始新的 Leader 臂校准 ---")
        try:
            self.bus_.disable_torque()
            for motor_name in self.joint_names:
                # 设置为位置模式
                self.bus_.write("Operating_Mode", motor_name, OperatingMode.POSITION.value)
            self.get_logger().info("所有电机已设置为位置模式。")

            # 1. 捕获 Homing Offsets
            input(">>> (1/2) 请用手将 Leader 臂移动到其运动范围的【中间位置】, 然后按 ENTER ...")
            homing_offsets = self.bus_.set_half_turn_homings()
            self.get_logger().info(f"成功捕获 Homing Offsets: {homing_offsets}")

            # 2. 捕获 Ranges
            full_turn_motor = "5"
            unknown_range_motors = [m for m in self.joint_names if m != full_turn_motor]

            print("\n>>> (2/2) 现在, 请【依次】将【除'5'号(手腕旋转)之外】的其它关节移动到它们的【最小值和最大值】。完成后按 ENTER 停止记录...\n")
            
            # 假设 bus 接口与 SO100Leader 示例一致
            range_mins, range_maxes = self.bus_.record_ranges_of_motion(unknown_range_motors)
            
            # 检查用户是否已经校准了 5
            if full_turn_motor not in range_mins:
                self.get_logger().info(f"未检测到关节 '5' 的范围, 假定为全周 (0-4095)。")
                range_mins[full_turn_motor] = 0
                range_maxes[full_turn_motor] = 4095
            else:
                self.get_logger().info(f"关节 '5' 已被校准, 使用记录的范围。")

            
            self.get_logger().info(f"记录 Range Mins: {range_mins}")
            self.get_logger().info(f"记录 Range Maxes: {range_maxes}")

            # 3. 构建完整的校准数据
            full_calibration_data = {}
            for motor_name, m in self.bus_.motors.items():
                if motor_name not in homing_offsets or motor_name not in range_mins:
                    self.get_logger().error(f"电机 {motor_name} 缺少校准数据！")
                    continue
                
                full_calibration_data[motor_name] = MotorCalibration(
                    id=m.id,
                    drive_mode=0, # Hardcoded from SO100Leader example
                    homing_offset=homing_offsets[motor_name],
                    range_min=range_mins[motor_name],
                    range_max=range_maxes[motor_name],
                )

            # 4. 将校准数据写入电机固件
            self.get_logger().info("正在将完整校准数据写入电机固件...")
            self.bus_.write_calibration(full_calibration_data) 
            self.get_logger().info("校准数据已写入电机。")

            # 5. 保存到本地文件
            self._save_calibration(full_calibration_data)
            
            return full_calibration_data
            
        except Exception as e:
            self.get_logger().error(f"校准过程中出错: {e}")
            import traceback
            self.get_logger().error(traceback.format_exc())
            return None

    def _load_or_run_calibration(self) -> bool:
        """
        - 如果 'calibrate' 参数为 True, 强制运行交互式校准。
        - 否则 (默认), 自动加载校准文件。
        - 如果文件不存在, 打印错误并失败。
        """
        
        # 模式 1: 强制校准 (用于手动运行)
        if self.is_calibrate_mode:
            self.get_logger().info("--- 运行交互式校准 ---")
            self.calibration_data = self._run_calibration()
            if self.calibration_data is None:
                self.get_logger().error("交互式校准失败。")
                return False
            self.get_logger().info("交互式校准成功。")
            return True # 校准成功

        # 模式 2: 自动加载 (用于 Launch 文件)
        if CALIB_PATH.is_file():
            self.get_logger().info(f"找到校准文件: {CALIB_PATH}。将自动加载此文件。")
            try:
                with open(CALIB_PATH, 'r') as f:
                    loaded_data = json.load(f) # 这是一个 dict[str, dict]
                
                self.calibration_data = {}
                for name, data_dict in loaded_data.items():
                    if name in self.joint_names: # 只加载知道的关节
                        self.calibration_data[name] = MotorCalibration(**data_dict) # 实例化
                
                # 基础验证
                if (not all(j in self.calibration_data for j in self.joint_names)):
                     self.get_logger().error("校准文件无效或不完整。请删除旧文件并运行手动校准。")
                     self.get_logger().error("  ros2 run so101_hw_interface so101_leader_pub --ros-args -p calibrate:=True")
                     return False
                else:
                    self.get_logger().info("已从文件加载校准数据。")
                    self.get_logger().info("正在将加载的校准数据写入电机固件...")
                    self.bus_.write_calibration(self.calibration_data)
                    self.get_logger().info("校准数据已写入电机。")
                    return True # Success
            except Exception as e:
                self.get_logger().error(f"加载/写入校准文件失败 ({e})。")
                self.get_logger().error("请删除损坏的文件并运行手动校准:")
                self.get_logger().error("  ros2 run so101_hw_interface so101_leader_pub --ros-args -p calibrate:=True")
                return False # Failure
        
        # 模式 2 (续): 文件未找到
        else:
            self.get_logger().error(f"未找到校准文件: {CALIB_PATH}")
            self.get_logger().error("请在【单独的终端中】运行此节点以执行交互式校准:")
            self.get_logger().error("  ros2 run so101_hw_interface so101_leader_pub --ros-args -p calibrate:=True")
            self.get_logger().error("校准完成后, 重新启动 launch 文件。")
            return False # Failure

    def timer_callback(self):
        try:
            # 读取原始位置
            raw_positions = self.bus_.sync_read("Present_Position", normalize=False)

            if not raw_positions or len(raw_positions) != 6:
                self.get_logger().warn(f"读取舵机位置失败或数据不全. 得到: {raw_positions}", throttle_duration_sec=5.0)
                return

            positions_rad = []
            for n in self.joint_names:
                raw = raw_positions.get(n, 0)
                home = self._home_offsets.get(n, 0) # 现在这是来自校准的真实零点
                
                # 1. 计算真实弧度
                true_rad = (raw - home) * self._rad_per_step 
                
                # 2. 获取此关节的 "Home" 弧度偏移量
                reset_offset = self._reset_rad_offsets.get(n, 0.0)
                
                # 3. 计算相对弧度 (真实弧度 - 偏移量)
                relative_rad = true_rad - reset_offset
                
                # 4. 添加相对弧度
                positions_rad.append(relative_rad)

            # 5. 只发布 Leader 自己的状态
            joint_state_msg = JointState()
            joint_state_msg.header.stamp = self.get_clock().now().to_msg()
            joint_state_msg.name = self.joint_names
            joint_state_msg.position = positions_rad
            
            self.leader_state_pub.publish(joint_state_msg)

        except Exception as e:
            self.get_logger().error(f"舵机读取循环出错: {e}", throttle_duration_sec=5.0)
            import traceback
            self.get_logger().error(traceback.format_exc())

def main(args=None):
    rclpy.init(args=args)
    node = None # 预先声明
    try:
        node = LeaderArmPublisher()
        # 只有在 __init__ 成功 (即校准成功) 且 bus 有效时才 spin
        if node.bus_ is not None and rclpy.ok() and node._home_offsets is not None: 
            rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        if node and node.bus_ is not None:
            node.get_logger().info("正在断开总线连接...")
            # 在断开连接前禁用扭矩
            try:
                node.bus_.disable_torque()
                node.get_logger().info("电机扭矩已禁用。")
            except Exception as e:
                node.get_logger().warn(f"禁用扭矩失败: {e}")
            node.bus_.disconnect()
        if node and rclpy.ok():
            node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()